// Copyright (c) 2021 Presto Labs Pte. Ltd.
// Author: xguo

#include "coin2/strategy/umm/pricer.h"

#include <algorithm>

namespace xguo {

inline bool IsAggOrder(const OrderData& order) {
  bool is_agg = order.type == OrderData::AGGRESSIVE_BUY || order.type == OrderData::AGGRESSIVE_SELL;
  return order.success && is_agg;
}

UmmPricer::UmmPricer(const UmmPricerConfig& config)
  : config_(config), name_(config.name) {
  MakeFillStats();
  for (int n = 0; n < config.num_products; n++) {
    rate_.push_back(std::make_unique<MovingWindowData<RateEntry>>(config_.rate_window_size));
    bbo_.push_back(Bbo());
    pid_to_index_[config.product_ids[n]] = n;
  }
  ohlc_ = std::make_unique<OhlcCalculator>(config_.ohlc_window_size);
  trade_window_ = std::make_unique<MovingWindowData<Bbo>>(config_.ohlc_window_size);
  use_tag_ = !!getenv("USE_TAG");

  LOG(INFO) << "Create UmmPricer: " << name_ << ", "
            << config.DebugString() << ", use_tag: " << use_tag_;
}

void UmmPricer::UpdateByBbo(const Bbo& bbo) {
  CHECK_GT(bbo.timestamp, timestamp_);
  timestamp_ = bbo.timestamp;
  CHECK_GT(timestamp_, 100000);

  int idx = pid_to_index_.at(bbo.product_id);
  bbo_[idx].CopyAndIncSeqNo(bbo);

  if (is_trade(bbo)) {
    fill_stats_->UpdateByBbo(bbo);
    ohlc_->UpdateByBbo(bbo);
    trade_window_->Update(bbo);
  }

  if (is_bbo_ready(0) && is_bbo_ready(idx)) {
    rate_[idx]->Update(RateEntry(bbo.timestamp, bbo_[0].midp / bbo_[idx].midp));
  }
}

void UmmPricer::UpdateByFill(const Fill& fill) {
  // update by fill
  fill_stats_->UpdateByFill(fill);
}

void UmmPricer::UpdateByPos(double pos) {
  // update by pos
  fill_stats_->UpdateByPos(pos - config_.reserve);

  if (stop_loss_mode_ == 1 && fill_stats_->position() <= +1e-5) {
    stop_loss_mode_ = 0;
  }

  if (stop_loss_mode_ == -1 && fill_stats_->position() > -1e-5) {
    stop_loss_mode_ = 0;
  }
}

void UmmPricer::TakeProfitOrStopLoss() {
  auto* fill_buy = fill_stats_->lowest_fill_buy();
  auto* fill_sell = fill_stats_->highest_fill_sell();
  CHECK(!(fill_buy && fill_sell)) << DebugString();

  if (fill_buy) {
    const double buy_rtn = std::log(trade_bbo().bid0 / fill_buy->fill_price);
    if (buy_rtn > config_.take_profit_rate) {
      if (config_.use_agg) {
        // If current buy price is much higher, then sell immediately.
        MakeTag();
        set_tag("type", "agg_sell_take_profit");
        set_tag("rtn", buy_rtn);

        order_data_.type = OrderData::AGGRESSIVE_SELL;
        order_data_.success = true;
        order_data_.agg_sell_price = trade_bbo().bid0;
        order_data_.agg_sell_qty = config_.lot_size;
        order_data_.set_tag(tag());
      } else {
        // If current buy price is much higher, then sell at ask0
        MakeTag();
        set_tag("type", "pass_sell_take_profit");
        set_tag("rtn", buy_rtn);

        order_data_.type = OrderData::PASSIVE;
        order_data_.success = true;
        order_data_.pass_sell_price = trade_bbo().ask0;
        order_data_.pass_sell_qty = config_.lot_size;
        order_data_.set_tag(tag());
      }
      CHECK_EQ(fill_sell, nullptr) << DebugString();
    } else if (buy_rtn < config_.stop_loss_rate) {
      const auto dt = timestamp_ - fill_stats_->last_sell_fill_timestamp();

      if (dt > config_.max_nofill_window_size) {
        stop_loss_mode_ = 1;
      }

      if (stop_loss_mode_ != 0) {
        if (config_.use_agg) {
          // If current buy price is much higher, then sell immediately.
          MakeTag();
          set_tag("type", "agg_sell_stop_loss");
          set_tag("rtn", buy_rtn);

          order_data_.success = true;
          order_data_.type = OrderData::AGGRESSIVE_SELL;
          order_data_.agg_sell_price = trade_bbo().bid0;
          order_data_.agg_sell_qty = config_.lot_size;
          order_data_.set_tag(tag());
        } else {
          MakeTag();
          set_tag("type", "pass_sell_stop_loss");
          set_tag("rtn", buy_rtn);

          order_data_.success = true;
          order_data_.type = OrderData::PASSIVE;
          order_data_.pass_sell_price = trade_bbo().ask0;
          order_data_.pass_sell_qty = config_.lot_size;
          order_data_.set_tag(tag());
        }
        CHECK_EQ(fill_sell, nullptr) << DebugString();
      }
    }
  }

  // If current sell price is much lower, then buy immediately.
  if (fill_sell) {
    const double sell_rtn = std::log(fill_sell->fill_price / trade_bbo().ask0);
    if (sell_rtn > config_.take_profit_rate) {
      if (config_.use_agg) {
        MakeTag();
        set_tag("type", "agg_buy_take_profit");
        set_tag("rtn", sell_rtn);

        order_data_.type = OrderData::AGGRESSIVE_BUY;
        order_data_.success = true;
        order_data_.agg_buy_price = trade_bbo().ask0;
        order_data_.agg_buy_qty = config_.lot_size;
        order_data_.set_tag(tag());
      } else {
        MakeTag();
        set_tag("type", "pass_buy_take_profit");
        set_tag("rtn", sell_rtn);

        order_data_.type = OrderData::PASSIVE;
        order_data_.success = true;
        order_data_.pass_buy_price = trade_bbo().bid0;
        order_data_.pass_buy_qty = config_.lot_size;
        order_data_.set_tag(tag());
      }
      CHECK_EQ(fill_buy, nullptr) << DebugString();
    } else if (sell_rtn < config_.stop_loss_rate) {
      const auto dt = timestamp_ - fill_stats_->last_sell_fill_timestamp();

      if (dt > config_.max_nofill_window_size) {
        stop_loss_mode_ = -1;
        // LOG(ERROR) << "stop loss mode: " << stop_loss_mode_;
      }

      if (stop_loss_mode_ != 0) {
        if (config_.use_agg) {
          // If current sell price is much higher, then buy immediately.
          MakeTag();
          set_tag("type", "agg_buy_stop_loss");
          set_tag("rtn", sell_rtn);

          order_data_.type = OrderData::AGGRESSIVE_BUY;
          order_data_.success = true;
          order_data_.agg_buy_price = trade_bbo().ask0;
          order_data_.agg_buy_qty = config_.lot_size;
          order_data_.set_tag(tag());
        } else {
          // If current sell price is much higher, then buy immediately.
          MakeTag();
          set_tag("type", "pass_buy_stop_loss");
          set_tag("rtn", sell_rtn);

          order_data_.type = OrderData::PASSIVE;
          order_data_.success = true;
          order_data_.pass_buy_price = trade_bbo().bid0;
          order_data_.pass_buy_qty = config_.lot_size;
          order_data_.set_tag(tag());
        }
        CHECK_EQ(fill_buy, nullptr) << DebugString();
      }
    }
  }
}

const OrderData& UmmPricer::CalculateOrder() {
  order_data_.Clear();
  order_data_.ask0 = trade_bbo().ask0;
  order_data_.bid0 = trade_bbo().bid0;

  // if any ref ready
  if (!is_bbo_ready(0) || !is_any_ref_ready()) {
    order_data_.success = false;
    return order_data_;
  }

  if (!is_fire_cool()) {
    order_data_.success = false;
    return order_data_;
  }

  TakeProfitOrStopLoss();
  if (order_data_.success) {
    if (IsAggOrder(order_data_)) {
      if (is_agg_cool()) {
        last_agg_time_ = timestamp_;
      } else {
        order_data_.success = false;
      }
      return order_data_;
    } else if (order_data_.type == OrderData::PASSIVE) {
      return order_data_;
    }
  }

#if 0
  if (config_.force_on_level) {
    order_data_.type = OrderData::PASSIVE;
    order_data_.success = true;

    order_data_.pass_sell_price = trade_bbo().ask0;
    order_data_.pass_sell_qty = config_.lot_size;

    order_data_.pass_buy_price = trade_bbo().bid0;
    order_data_.pass_buy_qty = config_.lot_size;

    order_data_.set_tag("passive_force_on_level");
    return order_data_;
  }
#endif

#if 0

  // Calculate mm price based on historical volatility.
  ohlc_->Sort();
  double scale = std::sqrt(config_.order_update_period_sec);
  double trade_edge = ohlc_->volatility() * scale;
  double last_edge = ohlc_->last_volatility() * scale;

  // at least 2bp
  trade_edge = std::max({trade_edge, last_edge, config_.maker_fee_rate, 2e-4});

  double trade_buy_price = trade_window_->GetMedian() * (1 - trade_edge);
  double trade_sell_price = trade_window_->GetMedian() * (1 + trade_edge);

  // record debug info
  order_data_.trade_edge = trade_edge;
  order_data_.buy_price[0] = trade_buy_price;
  order_data_.sell_price[0] = trade_sell_price;
#endif
  order_data_.buy_price[0] = order_data_.bid0;
  order_data_.sell_price[0] = order_data_.ask0;

  const int nprod = config_.num_products;
  // Calculate reference price.
  for (int n = 1; n < nprod; n++) {
    if (rate_[n]->size() < 10) {
      continue;
    }
    double ref_buy_rate = rate_[n]->GetLow();
    double ref_sell_rate = rate_[n]->GetHigh();
    double ref_buy_price = bbo_[n].midp * ref_buy_rate;
    double ref_sell_price = bbo_[n].midp * ref_sell_rate;

    // record debug info
    order_data_.buy_price[n] = ref_buy_price;
    order_data_.sell_price[n] = ref_sell_price;
    order_data_.ref_buy_rate = ref_buy_rate;
    order_data_.ref_sell_rate = ref_sell_rate;
  }

  auto* fill_buy = fill_stats_->lowest_fill_buy();
  auto* fill_sell = fill_stats_->highest_fill_sell();
  if (fill_sell) {
    order_data_.sell_price[nprod] = fill_sell->fill_price * (1 + config_.walk_step);
    order_data_.buy_price[nprod] = fill_sell->fill_price * (1 - config_.walk_step);
  }
  if (fill_buy) {
    order_data_.sell_price[nprod] = fill_buy->fill_price * (1 + config_.walk_step);
    order_data_.buy_price[nprod] = fill_buy->fill_price * (1 - config_.walk_step);
  }
  CHECK(!(fill_buy && fill_sell));

  CalculatePassivePrice();

  order_data_.type = OrderData::PASSIVE;
  order_data_.success = true;
  order_data_.pass_sell_qty = config_.lot_size;
  order_data_.pass_buy_qty = config_.lot_size;

  order_data_.set_tag(tag());
  return order_data_;
}

void UmmPricer::CalculatePassivePrice() {
  const int nprod = config_.num_products + 1;
  MakeTag();
  set_tag("type", "mm_order");

  double sell_price = 0.0;
  for (int n = 0; n < nprod; n++) {
    double p = order_data_.sell_price[n];
    if (p > sell_price) {
      sell_price = p;
    }
  }

  double buy_price = sell_price;
  for (int n = 0; n < nprod; n++) {
    double p = order_data_.buy_price[n];
    if (p > 0 && p < buy_price) {
      buy_price = p;
    }
  }

  // FIXME(xguo) consider the case spread is big.
  order_data_.pass_buy_price = std::min(buy_price, order_data_.bid0);
  order_data_.pass_sell_price = std::max(sell_price, order_data_.ask0);
  CHECK_GT(order_data_.pass_buy_price, 0) << order_data_.DebugString();
  CHECK_GT(order_data_.pass_sell_price, 0) << order_data_.DebugString();;
  set_tag("pass_buy_price", order_data_.pass_buy_price);
  set_tag("pass_sell_price", order_data_.pass_sell_price);
}

std::string UmmPricer::DebugString() const {
  coin2::strategy::umm::UmmPricerProto proto;

  auto fs = proto.mutable_fill_stats();
  fs->MergeFrom(fill_stats_->DebugProto());

  /*
  for (const auto& ohlc : ohlc_->IterOhlc()) {
    auto* entry = proto.add_ohlc();
    entry->set_timestamp(ohlc.timestamp);
    entry->set_open(ohlc.open);
    entry->set_high(ohlc.high);
    entry->set_low(ohlc.low);
    entry->set_close(ohlc.close);
    entry->set_rtn(ohlc.rtn);
  }
  */

  return proto.DebugString();
}

}  // namespace xguo
